# Datastructures defining an input batch

from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Set

import numpy as np
import torch

from vllm.multimodal         import MultiModalKwargs
from vllm.sampling_params    import SamplingParams, SamplingType
from vllm.v1.sample.metadata import SamplingMetadata

if TYPE_CHECKING:
    from vllm.multimodal.inputs import PlaceholderRange

#! 缓存的请求状态
# 已经提交到vLLM的请求可能很多，有的已经生成了部分token，有的可能还在排队等，因此要进行记录
@dataclass
class CachedRequestState:
    req_id: str                                      # request的id
    prompt_token_ids: List[int]                      # 该请求的输入内容的token id
    prompt: Optional[str]                            # 该请求的输入内容
    mm_inputs: List[MultiModalKwargs]                # 多模态模型的输入参数信息
    mm_positions: List["PlaceholderRange"]           # 多模态数据的位置和偏移信息
    sampling_params: SamplingParams                  # 输出时的采样参数
    generator: Optional[torch.Generator]             #! 用于管理伪随机数生成器状态的类
    # torch.rand(2,3, generator=generator)           # 通过指定的生成器生成随机数

    block_ids: List[int]                             # TODO: 使用了那些块？
    num_computed_tokens: int                         # 已经生成完成的token数量
    output_token_ids: List[int]                      # 输出token的内容

    @property
    def num_tokens(self) -> int:                     # 返回该请求的总token数量
        return len(self.prompt_token_ids) + len(self.output_token_ids)


class InputBatch:
    """用于batch推理的输入batch表达的类
    # TODO(yangxianpku): 没有beam search的设置
    """
    def __init__(
        self,
        max_num_reqs:  int,                  #! 一个batch中的最大请求数量
        max_model_len: int,                  # 输入序列的最大长度
        max_num_blocks_per_req: int,         # TODO: 每个请求占用的最大块数量
        device: torch.device,
        pin_memory: bool,                    # TODO: 昇腾的设备是否存在pinned的概念
        vocab_size: int,                     # TODO: 词汇表大小
    ):
        self.device        = device
        self.pin_memory    = pin_memory
        self.vocab_size    = vocab_size
        self.max_num_reqs  = max_num_reqs
        self.max_model_len = max_model_len
        self.max_num_blocks_per_req = max_num_blocks_per_req

        self.req_ids: List[Optional[str]]    = [None] * max_num_reqs
        self.req_id_to_index: Dict[str, int] = {}   #!  str类型的请求id到int类型的id映射关系

        # TODO(woosuk): This buffer could be too large if max_model_len is big.
        # Find a way to reduce the CPU memory usage.
        # woosuk应该是一个作者的名字
        self.token_ids_cpu_tensor = torch.zeros(
            (max_num_reqs, max_model_len),
            device="cpu",
            dtype=torch.int32,
            pin_memory=pin_memory,
        )
        self.token_ids_cpu           = self.token_ids_cpu_tensor.numpy()        # TODO(yangxian)  why ？
        self.num_computed_tokens_cpu = np.empty(max_num_reqs, dtype=np.int32)
        self.num_prompt_tokens       = np.zeros(max_num_reqs, dtype=np.int32)

        # Attention-related.
        self.block_table = torch.zeros(                #! TODO: 块表，用于管理逻辑block与物理block之间的映射
            (max_num_reqs, max_num_blocks_per_req),
            device=self.device,
            dtype=torch.int32,
        )
        self.block_table_cpu_tensor = torch.zeros(                 # TODO: # TODO(yangxian)  why + 1？
            (max_num_reqs, max_num_blocks_per_req),
            device="cpu",
            dtype=torch.int32,
            pin_memory=pin_memory,
        )
        self.block_table_cpu = self.block_table_cpu_tensor.numpy() # TODO: # TODO(yangxian)  why +2 ？

        # Sampling-related.
        #! 采样细节讨论： https://zhuanlan.zhihu.com/p/13873178220
        #! torch.empty 快速分配存储空间但不初始化数据，数据内容是随机的，但性能比torch.zeros要高
        self.temperature = torch.empty((max_num_reqs, ),           #! 每个请求对应一个采样的温度系数
                                       dtype=torch.float32,
                                       device=device)
        self.temperature_cpu_tensor = torch.empty((max_num_reqs, ),
                                                  dtype=torch.float32,
                                                  device="cpu",
                                                  pin_memory=pin_memory)
        self.temperature_cpu = self.temperature_cpu_tensor.numpy()
        self.greedy_reqs: Set[str] = set()      #! 贪心采样的请求（每次采样概率最大的token, 容易出现重复回答）
        self.random_reqs: Set[str] = set()      # TODO: 做什么用？  随机采样的请求？

        #! top-p采样，也叫核采样(Nucleus Sampling)
        # 入围概率值超过阈值top-p的n个token, 重新进行归一化保证总概率为1，再根据归一化后的概率随机选择一个token
        self.top_p = torch.empty((max_num_reqs, ),
                                 dtype=torch.float32,
                                 device=device)
        self.top_p_cpu_tensor = torch.empty((max_num_reqs, ),
                                            dtype=torch.float32,
                                            device="cpu",
                                            pin_memory=pin_memory)
        self.top_p_cpu = self.top_p_cpu_tensor.numpy()
        self.top_p_reqs: Set[str] = set()        #! 使用top_p策略采样的请求

        #! 找到概率最高的前k的token, 重新进行归一化保证总概率为1，再根据归一化后的概率随机选择一个token
        #! top_k + top_p 策略先试用top_k,再使用top_p
        # TODO: top_k后的概率是否需要归一化
        self.top_k = torch.empty((max_num_reqs, ),
                                 dtype=torch.int32,
                                 device=device)
        self.top_k_cpu_tensor = torch.empty((max_num_reqs, ),
                                            dtype=torch.int32,
                                            device="cpu",
                                            pin_memory=pin_memory)
        self.top_k_cpu = self.top_k_cpu_tensor.numpy()
        self.top_k_reqs: Set[str] = set()        #! 使用top_k策略采样的请求

        # Frequency penalty related data structures
        #! 1. 频率惩罚: 在logits中减去一定的值然后再softmax
        """
        frequency_penalty 是一个用于控制文本生成多样性的参数。它基于新生成的标记在当前生成文本中的出现频率来施加惩罚
        
        作用
            - 减少重复性：通过惩罚已经频繁出现的标记，frequency_penalty 可以减少模型在生成文本时的重复性，从而提高文本的多样性。
            - 鼓励使用新标记：正值的 frequency_penalty 会鼓励模型使用新的标记，而不是重复已经使用过的标记。
        工作原理
            当某个标记在生成的文本中出现得越频繁，它的惩罚就越大，从而降低该标记再次被选中的概率。
            这有助于避免生成过于冗长或单调的文本。
            取值范围
            frequency_penalty 的取值范围是 −2.0,2.0，默认值为 0.0。
            值 > 0：鼓励模型使用新标记，减少重复。
            值 < 0：鼓励模型重复使用标记，增加重复性。
        """
        self.frequency_penalties = torch.empty((max_num_reqs, ),
                                               dtype=torch.float,
                                               device=device)
        self.frequency_penalties_cpu_tensor = torch.empty(
                                                (max_num_reqs, ),
                                                dtype=torch.float,
                                                device="cpu",
                                                pin_memory=pin_memory)
        self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy()
        self.frequency_penalties_reqs: Set[str] = set()   #! 1. 施加频率惩罚的请求

        # Presence penalty related data structures
        #! 2. 存在性惩罚：在logits中减去一定的值然后再softmax
        self.presence_penalties = torch.empty((max_num_reqs, ),
                                              dtype=torch.float,
                                              device=device)
        self.presence_penalties_cpu_tensor = torch.empty((max_num_reqs, ),
                                                         dtype=torch.float,
                                                         device="cpu",
                                                         pin_memory=pin_memory)
        self.presence_penalties_cpu = \
            self.presence_penalties_cpu_tensor.numpy()
        self.presence_penalties_reqs: Set[str] = set()   #! 2. 施加存在性惩罚的请求

        # Repetition penalty related data structures
        #! 3. 重复惩罚：对logits缩放一定的值然后再softmax
        self.repetition_penalties = torch.empty((max_num_reqs, ),
                                                dtype=torch.float,
                                                device=device)
        self.repetition_penalties_cpu_tensor = torch.empty(
            (max_num_reqs, ),
            dtype=torch.float,
            device="cpu",
            pin_memory=pin_memory)
        self.repetition_penalties_cpu = \
            self.repetition_penalties_cpu_tensor.numpy()
        self.repetition_penalties_reqs: Set[str] = set()  #! 3. 施加重复惩罚的请求

        self.min_tokens: List[int] = [0] * max_num_reqs   #! 最小token id
        self.stop_token_ids: List[Set[int]] = [           # 每个请求的停止token
            set() for _ in range(max_num_reqs)            #! 模型启动时或读取停止token, 但也允许用户在请求中单独设置停止token
        ]

        # TODO: 为什么只有一个？
        self.prompt_token_ids: Optional[torch.Tensor] = None

        # req_index -> generator
        # NOTE(woosuk): The indices of the requests that do not have their own
        # generator should not be included in the dictionary.
        # TODO(yangxian):  不拥有独立生成器的请求不用放在该dict中？
        self.generators: Dict[int, torch.Generator] = {}  #! 每个请求对应一个随机数的生成器

        #! logprobs 是每个预测token的预测概率
        #  1. 由于概率值通常很小，直接使用概率可能会导致数值下溢（underflow），因此使用对数概率来避免这一问题。
        #  2. 对数概率可以将多个概率的乘法转换为加法，从而简化计算
        self.num_logprobs: Dict[str, int]  = {}           #! 生成每个toekn时，返回logprobs个最可能得token的对数概率
        self.prompt_logprob_reqs: Set[str] = set()        #! 每个输入提示词上返回对数概率的数量

    #! 向输入中添加请求
    def add_request(
        self,
        request: "CachedRequestState",
        req_index: Optional[int] = None,
    ) -> None:
        if req_index is None:
            req_index = self.num_reqs
        assert req_index < self.max_num_reqs             # TODO: 这样做是否合适？

        req_id                       = request.req_id    # str
        self.req_ids[req_index]      = req_id
        self.req_id_to_index[req_id] = req_index

        # Copy the prompt token ids and output token ids.
        num_prompt_tokens = len(request.prompt_token_ids)
        self.num_prompt_tokens[req_index] = num_prompt_tokens
        self.token_ids_cpu[
            req_index, :num_prompt_tokens] = request.prompt_token_ids
        start_idx = num_prompt_tokens
        end_idx   = start_idx + len(request.output_token_ids)           # TODO: 输出的token已经有了？
        self.token_ids_cpu[req_index,
                           start_idx:end_idx] = request.output_token_ids

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
        num_blocks = len(request.block_ids)
        self.block_table_cpu[req_index, :num_blocks] = request.block_ids # TODO: 使用了哪些块

        sampling_params = request.sampling_params
        self.temperature_cpu[req_index] = sampling_params.temperature
        if sampling_params.sampling_type == SamplingType.GREEDY:
            self.greedy_reqs.add(req_id)
        else:
            self.random_reqs.add(req_id)

        self.top_p_cpu[req_index] = sampling_params.top_p
        if sampling_params.top_p < 1:
            self.top_p_reqs.add(req_id)
        self.top_k_cpu[req_index] = sampling_params.top_k
        if sampling_params.top_k > 0:
            self.top_k_reqs.add(req_id)
        self.frequency_penalties_cpu[req_index] = \
            sampling_params.frequency_penalty
        if sampling_params.frequency_penalty != 0.0:     #! 直接添加请求id (str)
            self.frequency_penalties_reqs.add(req_id)
        self.presence_penalties_cpu[req_index] = \
            sampling_params.presence_penalty
        if sampling_params.presence_penalty != 0.0:
            self.presence_penalties_reqs.add(req_id)
        self.repetition_penalties_cpu[req_index] = \
            sampling_params.repetition_penalty
        if sampling_params.repetition_penalty != 1.0:
            self.repetition_penalties_reqs.add(req_id)
        self.min_tokens[req_index]     = sampling_params.min_tokens
        self.stop_token_ids[req_index] = sampling_params.all_stop_token_ids

        # NOTE(woosuk): self.generators should not include the requests that
        # do not have their own generator.
        if request.generator is not None:
            self.generators[req_index] = request.generator

        num_logprobs = sampling_params.logprobs
        if num_logprobs is not None and num_logprobs > 0:
            self.num_logprobs[req_id] = num_logprobs
        if sampling_params.prompt_logprobs:
            self.prompt_logprob_reqs.add(req_id)

    #! 某个请求处理完成了，可以移除
    def remove_request(self, req_id: str) -> Optional[int]:
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
        self.req_ids[req_index] = None

        self.greedy_reqs.discard(req_id)   #! discard 从集合中移除元素
        self.random_reqs.discard(req_id)
        self.top_p_reqs.discard(req_id)
        self.top_k_reqs.discard(req_id)
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
        self.prompt_logprob_reqs.discard(req_id)
        return req_index

    #! 全部清除
    def clear(self) -> None:
        self.req_ids = [None] * self.max_num_reqs
        self.req_id_to_index.clear()
        self.greedy_reqs.clear()
        self.random_reqs.clear()
        self.top_p_reqs.clear()
        self.top_k_reqs.clear()
        self.frequency_penalties_reqs.clear()
        self.presence_penalties_reqs.clear()
        self.repetition_penalties_reqs.clear()
        self.generators.clear()
        self.num_logprobs.clear()
        self.prompt_logprob_reqs.clear()

    #! 压缩：
    # TODO: 这块没搞明白，从Python是线上看有较大的优化空间
    def condense(self, empty_req_indices: List[int]) -> None:
        if self.num_reqs == 0:
            # The batched states are empty.
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order. 假设空请求时降序排列的
        # TODO: 什么是empty req？
        last_req_index = self.num_reqs + len(empty_req_indices) - 1
        while empty_req_indices:
            # Find the largest non-empty index.
            # 找到最大的非空索引
            while last_req_index in empty_req_indices:
                last_req_index -= 1

            # Find the smallest empty index.
            empty_index = empty_req_indices.pop()
            if empty_index >= last_req_index:
                break

            # Swap the states.
            req_id = self.req_ids[last_req_index]   # List[str]
            assert req_id is not None
            self.req_ids[empty_index] = req_id
            self.req_ids[last_req_index] = None
            self.req_id_to_index[req_id] = empty_index

            # TODO(woosuk): Optimize the copy of token_ids_cpu and
            # block_table_cpu.
            self.token_ids_cpu[empty_index] = self.token_ids_cpu[
                last_req_index]
            self.num_prompt_tokens[empty_index] = \
                self.num_prompt_tokens[last_req_index]
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
            self.block_table_cpu[empty_index] = self.block_table_cpu[
                last_req_index]
            self.temperature_cpu[empty_index] = self.temperature_cpu[
                last_req_index]
            self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index]
            self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index]
            self.frequency_penalties_cpu[empty_index] = \
                self.frequency_penalties_cpu[last_req_index]
            self.presence_penalties_cpu[empty_index] = \
                self.presence_penalties_cpu[last_req_index]
            self.repetition_penalties_cpu[empty_index] = \
                self.repetition_penalties_cpu[last_req_index]
            self.min_tokens[empty_index] = self.min_tokens[last_req_index]
            self.stop_token_ids[empty_index] = \
                self.stop_token_ids[last_req_index]
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

    #! 生成该请求批次的解码输出采样参数对象
    def make_sampling_metadata(
        self,
        req_id_output_token_ids: Dict[str, List[int]],
        skip_copy: bool = False,
    ) -> SamplingMetadata:            #! 生成采样参数的元数据dataclass
        if not skip_copy:
            self.temperature[:self.num_reqs].copy_(
                self.temperature_cpu_tensor[:self.num_reqs], non_blocking=True)
            self.top_p[:self.num_reqs].copy_(
                self.top_p_cpu_tensor[:self.num_reqs], non_blocking=True)
            self.top_k[:self.num_reqs].copy_(
                self.top_k_cpu_tensor[:self.num_reqs], non_blocking=True)
            if not self.no_penalties:
                # Since syncing these tensors is expensive only copy them
                # if necessary i.e. if there are requests which require
                # penalties to be applied during sampling.
                self.frequency_penalties[:self.num_reqs].copy_(
                    self.frequency_penalties_cpu_tensor[:self.num_reqs],
                    non_blocking=True)
                self.presence_penalties[:self.num_reqs].copy_(
                    self.presence_penalties_cpu_tensor[:self.num_reqs],
                    non_blocking=True)
                self.repetition_penalties[:self.num_reqs].copy_(
                    self.repetition_penalties_cpu_tensor[:self.num_reqs],
                    non_blocking=True)
                # The prompt tokens are used only for applying penalties during
                # the sampling process. Hence copy these tensors only when
                # there are requests which need penalties to be applied.
                self.prompt_token_ids = self._make_prompt_token_ids_tensor()

        output_token_ids: List[List[int]] = []

        for req_id in self.req_ids[:self.num_reqs]:
            assert req_id is not None
            # Currently we create a tensor for output_token_ids from scratch
            # at each step. However, for the penalties computation what we
            # need is stats about the token ids present in the output. This
            # stats can be maintained incrementally instead of computing it
            # from scratch at each step.
            # TODO - Replace this with incremental update to output token
            # statistics.
            output_token_ids.append(req_id_output_token_ids[req_id])

        return SamplingMetadata(
            temperature=self.temperature[:self.num_reqs],
            all_greedy=self.all_greedy,
            all_random=self.all_random,
            top_p=self.top_p[:self.num_reqs],
            top_k=self.top_k[:self.num_reqs],
            no_top_p=self.no_top_p,
            no_top_k=self.no_top_k,
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
            prompt_token_ids=self.prompt_token_ids,
            frequency_penalties=self.frequency_penalties[:self.num_reqs],
            presence_penalties=self.presence_penalties[:self.num_reqs],
            repetition_penalties=self.repetition_penalties[:self.num_reqs],
            output_token_ids=output_token_ids,
            min_tokens=self.min_tokens[:self.num_reqs],
            stop_token_ids=self.stop_token_ids[:self.num_reqs],
            no_penalties=self.no_penalties,
        )

    #! 将一个batch中的请求的prompt组成一个单独的device tensor
    def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
        max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()  # 请求中最大的长度
        prompt_token_ids_cpu_tensor = torch.empty(
            (self.num_reqs, max_prompt_len),
            device="cpu",
            dtype=torch.int64,
            pin_memory=self.pin_memory)
        #! 预分配
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
        prompt_token_ids[:] = (
            self.token_ids_cpu[:self.num_reqs, :max_prompt_len])
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
        #! Note: 在一个batch中进行右侧补齐？ 左侧padding呢s？
        for i in range(self.num_reqs):
            prompt_token_ids[i, self.num_prompt_tokens[i]:] = self.vocab_size
        return prompt_token_ids_cpu_tensor.to(device=self.device,
                                              non_blocking=True)  #! 

    @property
    def num_reqs(self) -> int:
        """该batch中总的请求数量"""
        return len(self.req_id_to_index)

    @property
    def all_greedy(self) -> bool:
        """是否全部是贪心请求
        """
        return len(self.random_reqs) == 0

    @property
    def all_random(self) -> bool:
        """是否全部是随机请求: top-k, top-p等都属于随机请求
        """
        return len(self.greedy_reqs) == 0

    @property
    def no_top_p(self) -> bool:
        return len(self.top_p_reqs) == 0

    @property
    def no_top_k(self) -> bool:
        return len(self.top_k_reqs) == 0

    @property
    def no_penalties(self) -> bool:
        return (len(self.presence_penalties_reqs) == 0
                and len(self.frequency_penalties_reqs) == 0
                and len(self.repetition_penalties_reqs) == 0)

    @property
    def max_num_logprobs(self) -> int:
        return max(self.num_logprobs.values()) if self.num_logprobs else 0

    @property
    def no_logprob(self) -> bool:
        return len(self.num_logprobs) == 0

    @property
    def no_prompt_logprob(self) -> bool:
        return len(self.prompt_logprob_reqs) == 0
